import socket

class QwenCompatibleClient:
    def __init__(self, host='127.0.0.1', port=12345):
        self.host = host
        self.port = port
        self.sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
        self.sock.connect((self.host, self.port))
        print(f"[Connected] {self.host}:{self.port}")

    def _recv_until(self, delimiter=b"<END>"):
        buffer = b""
        while delimiter not in buffer:
            chunk = self.sock.recv(1024)
            if not chunk:
                break
            buffer += chunk
        return buffer.decode('utf-8', errors='ignore')

    def _extract_reply(self, full_text):
        start_tag = "<|im_start|>assistant"
        end_tag = "<END>"
        start_idx = full_text.find(start_tag)
        end_idx = full_text.find(end_tag, start_idx)
        if start_idx == -1 or end_idx == -1:
            return full_text.strip()
        return full_text[start_idx + len(start_tag):end_idx].strip()

    class chat:
        class completions:
            @staticmethod
            def create(model, messages, stream=False):
                # 直接合成一句话，忽略 OpenAI messages 复杂结构
                # 支持 messages 是 dict 或 list
                if isinstance(messages, dict):
                    system = messages.get("system", "")
                    user = messages.get("user", "")
                elif isinstance(messages, list):
                    system = ""
                    user = ""
                    for m in messages:
                        if m.get("role") == "system":
                            system = m.get("content", "")
                        elif m.get("role") == "user":
                            user = m.get("content", "")
                else:
                    system = ""
                    user = ""

                prompt = f"{system} {user}".strip()

                # 发送并接收
                QwenCompatibleClient._instance.sock.sendall((prompt + "\n").encode("utf-8"))
                full_text = QwenCompatibleClient._instance._recv_until()
                reply = QwenCompatibleClient._instance._extract_reply(full_text)

                return QwenResponse(reply)

    def close(self):
        try:
            self.sock.sendall(b"exit\n")
        except:
            pass
        self.sock.close()
        print("[Connection closed]")

    @classmethod
    def with_openai_api(cls, host='127.0.0.1', port=12345):
        inst = cls(host, port)
        cls._instance = inst
        return inst


class QwenResponse:
    class Message:
        def __init__(self, content):
            self.role = "assistant"
            self.content = content

    class Choice:
        def __init__(self, message):
            self.message = message

    def __init__(self, content):
        self.choices = [self.Choice(self.Message(content))]
